package demo1;

import java.util.*;

public class TopK {
    public static int[] topK (int[] array, int k) {
        //1.创建一个大小为K的大根堆
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>(k, new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o2 - o1;
            }
        });
        //2.遍历数组当中的元素,前K个元素放到队列当中
        for (int i = 0; i < array.length; i++) {
            if (maxHeap.size() < k) {
                maxHeap.offer(array[i]);
            } else {
                //3.从第K+1个元素开始,每个元素和堆顶元素进行比较
                int top = maxHeap.peek();
                if (top > array[i]) {
                    maxHeap.poll();
                    maxHeap.offer(array[i]);
                }
            }
        }
        int[] tmp = new int[k];
        for (int i = 0; i < k; i++) {
            tmp[i] = maxHeap.poll();
        }
        return tmp;
    }

    //找出前k对数字
    public static List<List<Integer>> kSamllestPairs (int[] nums1, int[] nums2,int k) {
        //建立一个大小为K的大根堆
        PriorityQueue<List<Integer>> maxHeap = new PriorityQueue<>(k, new Comparator<List<Integer>>() {
            @Override
            public int compare(List<Integer> o1, List<Integer> o2) {
                return (o2.get(0) + o2.get(1)) - (o1.get(0) + o1.get(1));
            }
        });
        //不用把nums1.length都走完   i < Math.min(nums1.length,k) 走k次就好了
        for (int i = 0; i < Math.min(nums1.length,k); i++) {
            for (int j = 0; j < Math.min(nums2.length,k); j++) {
                List<Integer> list = new ArrayList<>();
                list.add(nums1[i]);
                list.add(nums2[j]);
                if (maxHeap.size() < k) {
                    maxHeap.offer(list);
                } else {
                    List<Integer> list1 = maxHeap.peek();
                    if ((list1.get(0) + list1.get(1)) > (nums1[i] + nums2[j])) {
                        maxHeap.poll();
                        maxHeap.offer(list);
                    }
                }
            }
        }
        List<List<Integer>> ret = new ArrayList<>();
        int size = maxHeap.size();
        for (int i = 0; i < size; i++) {
            ret.add(maxHeap.poll());
        }
        return ret;
    }

    public static void main(String[] args) {
//        int[] array = new int[] {1,2,3,4,578,6,7,8,9,10};
//        int[] temp = topK(array, 3);
//        System.out.println(Arrays.toString(temp));
        int[] nums1 = {1,7,11};
        int[] nusm2 = {2,4,6};
        List<List<Integer>> ret = kSamllestPairs(nums1,nusm2,3);
        for (int i = 0; i < 2; i++) {
            System.out.println(ret.get(i).toString());
        }

    }
}
